import os
import os.path as osp
import random

BASEDIR = osp.abspath('.')
IMGDIR = osp.join(BASEDIR, 'img')
TXTDIR = osp.join(BASEDIR, 'new_txt_fatter')

OUTDIR = osp.join(BASEDIR, 'selected_train_val')
TRAIN_IMG_DIR = osp.join(OUTDIR, 'images', 'train')
TRAIN_TXT_DIR = osp.join(OUTDIR, 'labels', 'train')
VAL_IMG_DIR = osp.join(OUTDIR, 'images', 'val')
VAL_TXT_DIR = osp.join(OUTDIR, 'labels', 'val')

for DIR in [OUTDIR, TRAIN_IMG_DIR, TRAIN_TXT_DIR, VAL_IMG_DIR, VAL_TXT_DIR]:
    if not osp.exists(DIR):
        os.makedirs(DIR)

names = [name[:-4] for name in os.listdir(IMGDIR) if name.endswith('.tif')]
# print(names)
# print(len(names))

random.shuffle(names)
print(names)

for i in range(len(names)):
    tif_name = names[i]+'.tif'
    txt_name = names[i]+'.txt'

    src_path_tif = osp.abspath(osp.join(IMGDIR, tif_name))
    src_path_txt = osp.abspath(osp.join(TXTDIR, txt_name))

    if i < 60:  # val
        des_path_tif = osp.abspath(osp.join(VAL_IMG_DIR, tif_name))
        des_path_txt = osp.abspath(osp.join(VAL_TXT_DIR, txt_name))
    else:
        des_path_tif = osp.abspath(osp.join(TRAIN_IMG_DIR, tif_name))
        des_path_txt = osp.abspath(osp.join(TRAIN_TXT_DIR, txt_name))
        
    os.system(f'cp {src_path_tif} {des_path_tif}')
    os.system(f'cp {src_path_txt} {des_path_txt}')
    
    print(f'{names[i]} copied, done.')

